#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: zyx
@Date: 2024/12/08 15:33
@FileName: memory.py
@Description: langchain 记忆类型
"""

import os
import json
from typing import List, Optional
from collections.abc import Sequence
from dotenv import load_dotenv
from prompts import SUMMARY_TEMPLATE
from loguru import logger

from langchain_openai import ChatOpenAI
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage, AIMessage, HumanMessage
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.prompts import ChatPromptTemplate

from langchain.schema import messages_from_dict, message_to_dict

from langchain_community.utilities.redis import get_client


def summary_chat(store_message):
    llm = ChatOpenAI(model=os.environ["BASE_CHAT_LLM"], temperature=0)
    str_message = ""
    for message in store_message:
        str_message += f"{type(message).__name__}: {message.content}"
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", SUMMARY_TEMPLATE),
            ("user", "{input}"),
        ]
    )
    chain = prompt | llm
    summary = chain.invoke({"input": str_message})
    return summary


class InMemoryChatSummaryHistory(BaseChatMessageHistory):

    store: dict = {}

    def __init__(
        self,
        session_id: str,
        # trunc_len: 指定需要对历史信息做截断的长度，避免历史信息过长，需要为2的倍数（保证用户输入时进行总结）
        # trunc_len=0，表示不进行截断
        trunc_len: int = 10,
    ):
        self.session_id = session_id
        self.trunc_len = trunc_len

    @property
    def messages(self) -> List[BaseMessage]:
        if not self.session_id in InMemoryChatSummaryHistory.store:
            InMemoryChatSummaryHistory.store[self.session_id] = []
        return InMemoryChatSummaryHistory.store[self.session_id]

    @messages.setter
    def messages(self, messages: List[BaseMessage]) -> None:
        raise NotImplementedError(
            "Direct assignment to 'messages' is not allowed."
            " Use the 'add_messages' instead."
        )

    def add_message(self, message: BaseMessage) -> None:
        store_message = self.messages
        if (self.trunc_len == 0) or (len(store_message) < self.trunc_len):
            store_message.append(message)
            return
        # 对超长的聊天记录进行摘要
        summary = summary_chat(store_message)
        msg = HumanMessage(
            content=f"我们之前的聊天记录摘要是：{summary}，我当前提问为：{message.content}"
        )
        store_message.clear()
        store_message.append(msg)

    def clear(self) -> None:
        self.messages = []


class RedisChatSummaryHistory(BaseChatMessageHistory):
    """Chat message history stored in a Redis database."""

    def __init__(
        self,
        session_id: str,
        url: str = "redis://localhost:6379/0",
        key_prefix: str = "message_store:",
        ttl: Optional[int] = None,
        # trunc_len: 指定需要对历史信息做截断的长度，避免历史信息过长，需要为2的倍数（保证用户输入时进行总结）
        # trunc_len=0，表示不进行截断
        trunc_len: int = 10,
    ):
        try:
            import redis
        except ImportError:
            raise ImportError(
                "Could not import redis python package. "
                "Please install it with `pip install redis`."
            )
        try:
            self.redis_client = get_client(redis_url=url, decode_responses=True)
        except redis.exceptions.ConnectionError as error:
            logger.error(error)
        load_dotenv()
        self.session_id = session_id
        self.key_prefix = key_prefix
        self.ttl = ttl
        if trunc_len < 0 or trunc_len % 2 != 0:
            raise ValueError(
                f"trunc_len must be a positive even number, but got {trunc_len}"
            )
        self.trunc_len = trunc_len

    @property
    def key(self) -> str:
        return self.key_prefix + self.session_id

    @property
    def messages(self) -> List[BaseMessage]:
        _items = self.redis_client.lrange(self.key, 0, -1)
        items = [json.loads(m) for m in _items[::-1]]
        messages = messages_from_dict(items)
        return messages

    @messages.setter
    def messages(self, messages: List[BaseMessage]) -> None:
        raise NotImplementedError(
            "Direct assignment to 'messages' is not allowed."
            " Use the 'add_messages' instead."
        )

    def add_message(self, message: BaseMessage) -> None:
        store_message = self.messages
        if (self.trunc_len == 0) or (len(store_message) < self.trunc_len):
            self._add_msg(message)
            return
        # 对超长的聊天记录进行摘要
        summary = summary_chat(store_message)
        msg = HumanMessage(
            content=f"我们之前的聊天记录摘要是：{summary}，我当前提问为：{message.content}"
        )
        self._add_msg(msg)

    def _add_msg(self, msg: BaseMessage) -> None:
        self.redis_client.lpush(self.key, json.dumps(message_to_dict(msg), ensure_ascii=False))
        if self.ttl:
            self.redis_client.expire(self.key, self.ttl)

    def clear(self) -> None:
        self.redis_client.delete(self.key)


def get_memory_by_type(memory_type: str):
    match memory_type:
        case "memory":
            return lambda session_id: InMemoryChatSummaryHistory(session_id=session_id)
        case "redis":
            return lambda session_id: RedisChatSummaryHistory(
                session_id=session_id, url=os.environ["REDIS_URL"], ttl=60 * 60 * 24
            )
        case _:
            raise ValueError(f"不支持的记忆类型: {memory_type}")
